import torch


def pad_or_crop(x: torch.Tensor, d: int, dim: int) -> torch.Tensor:
    dim = dim if dim >= 0 else dim + len(x.shape)
    if x.shape[dim] < d:  # pad
        pad = torch.zeros([*x.shape[:dim], d - x.shape[dim], *x.shape[dim + 1:]], dtype=x.dtype, device=x.device)
        return torch.cat([x, pad], dim=dim)
    elif x.shape[dim] > d:  # crop
        slices = tuple([slice(None)] * dim + [slice(d)] + [slice(None)] * (len(x.shape) - dim - 1))
        return x[slices]
    else:  # same
        return x


def align(r: torch.Tensor, x: torch.Tensor, dim: int) -> torch.Tensor:
    return pad_or_crop(x, r.shape[dim], dim)


def aligning_add(r: torch.Tensor, x: torch.Tensor, dim: int) -> torch.Tensor:
    return r + align(r, x, dim)


def aligning_mul(r: torch.Tensor, x: torch.Tensor, dim: int) -> torch.Tensor:
    return r * align(r, x, dim)
